'''
Segment time series and store the results as a dict

'''
import os
import numpy as np
import torch

from model import segment_module
from experiment import exp_factory
import options

args = options.Options().parse()

F = ['TRAIN', 'TEST']
for flag in F:
    segment_idx = {}
    dataset, dataloader = exp_factory.data_factory(config=args, flag=flag)
    print(len(dataset))
    for i in range(len(dataset)):
        x, label, info = dataset[i]
        x = x.numpy()
        merged_labels, num_components_list = segment_module.hierarchicalsegment(x)
        print(info, x.shape, merged_labels.shape)
        segment_idx[info] = merged_labels

    savename = os.path.join(args.exp_segment_savepath, '{}_{}.npy'.format(args.data_name, flag))
    np.save(savename, segment_idx)

